import akshare as ak
import pandas as pd
from tqdm import tqdm


def get_and_export_zygc(output_path: str = "a股主营构成数据.xlsx") -> None:
    """
    获取A股及新三板（含92开头）所有公司主营构成并导出到Excel表格

    :param output_path: 导出文件路径（默认当前目录）
    """
    try:
        # 获取股票列表（列名为'code'和'name'）
        stock_df = ak.stock_info_a_code_name()
        # stock_df = stock_df.head(10)
        print(f"获取到 {len(stock_df)} 家上市公司及挂牌企业列表")

        # 存储主营构成数据的列表
        zygc_data = []

        # 遍历每支股票获取主营构成（带进度条）
        for _, row in tqdm(stock_df.iterrows(), total=len(stock_df), desc="数据获取进度"):

            stock_code = row['code']  # 原始代码（不带前缀）
            stock_name = row['name']  # 股票名称

            # 添加市场前缀（SH/SZ/BJ/OC）
            # if stock_code.startswith('6'):
            #     symbol = f"SH{stock_code}"  # 上交所（主板/科创板）
            # elif stock_code.startswith(('0', '3')):
            #     symbol = f"SZ{stock_code}"  # 深交所（主板/创业板）
            # elif stock_code.startswith('8'):
            #     symbol = f"BJ{stock_code}"  # 北交所（83/87/88/89等）
            # elif stock_code.startswith(('43', '92')):  # 新增92开头支持
            #     symbol = f"OC{stock_code}"  # 新三板（老三板，43/92等开头）
            # else:
            #     print(f"\n警告：股票 {stock_code}({stock_name}) 无法识别市场类型（非6/0/3/8/43/92开头），跳过获取")
            #     continue

            try:
                # 调用主营构成接口（使用带前缀的symbol）
                zygc_df = ak.stock_zyjs_ths(symbol=stock_code)
                if not zygc_df.empty:
                    # 添加原始股票代码和名称信息（不带前缀）
                    # zygc_df['原始代码'] = stock_code
                    zygc_df['股票名称'] = stock_name
                    zygc_data.append(zygc_df)
            except Exception as e:
                print(f"\n警告：股票 {stock_code}({stock_name}) 获取失败，错误信息：{str(e)}")
                continue

        if not zygc_data:
            print("未获取到任何主营构成数据")
            return

        print(f"已获取 {len(zygc_data)} 条数据")

        # 合并所有数据
        final_df = pd.concat(zygc_data, ignore_index=True)

        # 调整列顺序（核心信息前置）
        columns = ['股票名称'] + [col for col in final_df.columns if col not in ['原始代码', '股票名称']]
        final_df = final_df[columns]

        print(f"开始导出")

        # 导出到Excel（带自动列宽）
        with pd.ExcelWriter(output_path, engine='openpyxl') as writer:
            final_df.to_excel(writer, index=False, sheet_name='主营构成数据')
            # 自动调整列宽
            worksheet = writer.sheets['主营构成数据']
            for idx, col in enumerate(final_df.columns):
                max_len = max(final_df[col].astype(str).map(len).max(), len(col)) + 2
                worksheet.column_dimensions[chr(65 + idx)].width = max_len

        print(f"\n数据已成功导出至：{output_path}")

    except Exception as e:
        print(f"程序执行失败，错误信息：{str(e)}")


if __name__ == "__main__":
    # 可修改输出路径（如：D:/股市数据/主营构成.xlsx）
    get_and_export_zygc()
